# run
# python ./demo.py > result.txt


from random import randint
from numpy.random import zipf
import threading
import time
import numpy as np
import random
TIME_SCALE = 10

SLEEP_INTERVAL = 0.001*TIME_SCALE
IN_CACHE_TIME = 0.02*TIME_SCALE
OUT_CACHE_TIME = 0.04*TIME_SCALE
NORMAL_SEND_INTERVAL = 0.003*TIME_SCALE
EVIL_SEND_INTERVAL = 0.0001*TIME_SCALE

GLOBAL_START = time.time()
# SIMULATION_TIME = 3
SIMULATION_TIME = 50

ZF_A = 1.5 # zipf的参数

stats = {
    "normal": 0,
    "polluted": 0,
    "defended": 0
}


class Packet():
    def __init__(self, content_id, consumer_id, consumer_type):
        self.content_id = content_id
        self.consumer_id = consumer_id
        self.consumer_type = consumer_type
        self.send_timestamp = time.time()
        self.recv_timestamp = GLOBAL_START
        self.elapsed_time = 0
        self.in_cache = -1
    def set_recv_timestamp(self):
        self.recv_timestamp = time.time()
        self.elapsed_time = self.recv_timestamp - self.send_timestamp
    def set_cache_states(self, in_cache:bool):
        self.in_cache = in_cache
        # self.elapsed_time = 0.02 if self.in_cache else 0.04
    def __str__(self):
        return f"[{self.content_id}, {self.consumer_id}, {self.consumer_type}, {self.in_cache}, "\
               f"{self.send_timestamp-GLOBAL_START}, {self.recv_timestamp-GLOBAL_START}, {self.elapsed_time/TIME_SCALE}]"


class Link():
    def __init__(self, capcity:int):
        self.capcity = capcity
        self.used = 0
    def is_sendable(self):
        return self.used < self.capcity
    def consumer2node(self):
        # to add lock
        self.used+=1
    def node2consumer(self):
        # to add lock
        self.used-=1

class Node():
    def __init__(self, link:Link, cache_size=100):
        self.link = link
        self.cache = [-1 for i in range(cache_size)]
        self.record = []
        self.record_cache = []
    
    def update_cache(self, content_id):
        """[summary]

        Args:
            content_id ([type]): [description]

        Returns:
            int: 1: in cache, 0 not in cache
        """
        if content_id in self.cache:
            index = self.cache.index(content_id)
            self.cache = [content_id] + self.cache[:index] + self.cache[index+1:]
            in_cache=True
        else:
            self.cache = [content_id] + self.cache[:-1]
            in_cache=False
        
        #print(f"Node: in_cache={in_cache}, cache updated: {self.cache}")
        return in_cache
        
    def recv(self, packet:Packet):
        # add stats
        if packet.consumer_id > 0:
            stats['normal'] += 1
        else:
            stats['polluted'] += 1
            
        self.link.consumer2node()
        in_cache = self.update_cache(packet.content_id)
        packet.set_cache_states(in_cache)
        self.record.append(packet)
        import copy
        self.record_cache.append(copy.deepcopy(self.cache))
        flight_time = IN_CACHE_TIME if in_cache else OUT_CACHE_TIME
        return flight_time
        
    def send(self, packet:Packet):
        #print(f"Node: packet send back to consumer: {packet}")
        packet.set_recv_timestamp()
        self.link.node2consumer()

        

class Consumer():
    def __init__(self, node:Node, consumer_id:int, consumer_type:int, send_interval:float, content_size=100, send_prob=1.0):
        # consumer_type 没用上
        self.node = node
        self.consumer_id = consumer_id
        self.consumer_type = consumer_type
        self.content_size = content_size
        self.send_interval = send_interval
        self.send_prob = send_prob
    
    def generate_packet(self):
        if self.consumer_id > 0:
            # 正常用户
            content_id = zipf(ZF_A, 1)[0]
        else:
            # 不正常用户
            content_id = 0-randint(100, 106)
        return Packet(content_id, self.consumer_id, self.consumer_type)

    def send(self):
        def call_recv(flight_time, packet):
            time.sleep(flight_time)
            self.node.send(packet)
            #print(f"Consumer[{self.consumer_id}]: recv packet back from node: {packet}")
        while self.node.link.is_sendable() == False:
            #print(f"Consumer[{self.consumer_id}] tries to send packet but failed due to capcity not enough")
            time.sleep(self.send_interval) # 发送失败重复尝试
        
        packet = self.generate_packet()
        ret = random.choices([True, False],weights=[self.send_prob, 1-self.send_prob])[0]
        if not ret:
            stats['defended'] += 1
            return
        #print(f"Consumer[{self.consumer_id}]: send packet to node: {packet}")
        flight_time = self.node.recv(packet)
        t = threading.Thread(target=call_recv, args=(flight_time, packet))
        t.setDaemon(True)
        t.start()
        
    def run(self):
        while True:
            if time.time() - GLOBAL_START > SIMULATION_TIME: return 
            time.sleep(self.send_interval) # 发包间隔
            self.send()

            
def test():
    link = Link(260)
    node = Node(link, 10)
    num_normal_consumer = 100
    num_evil_consumer = 10
    normal_consumers = []
    evil_consumers = []
    for i in range(num_normal_consumer):
        consumer_id = i+1
        normal_consumers.append(Consumer(node, consumer_id, 0, NORMAL_SEND_INTERVAL, 100, 1.0))
    for i in range(num_evil_consumer):
        consumer_id = -i-1
        evil_consumers.append(Consumer(node, consumer_id, 0, NORMAL_SEND_INTERVAL, 100, 1.0))
        # evil_consumers.append(Consumer(node, consumer_id, 0, NORMAL_SEND_INTERVAL, 100, 0.06))

    threads = []
    for consumer in normal_consumers:
        threads.append(threading.Thread(target=consumer.run))
    for consumer in evil_consumers:
        threads.append(threading.Thread(target=consumer.run))
    for t in threads:
        t.setDaemon(True)
        t.start()
    
    
    for t in threads:
        t.join()

    #print(f"record: {node.record}")
    #print(f"record:\n")
    # for packet in node.record:
    #     # print(f"{packet}")
    #     print(f"{packet.send_timestamp-GLOBAL_START}\t{packet.content_id}\t{packet.consumer_id}\t{packet.in_cache}")
    

    for packet, cache in zip(node.record, node.record_cache):
        # print(f"{packet.send_timestamp-GLOBAL_START},{packet.content_id},{packet.consumer_id},{packet.in_cache}\t")
        pass

    normal = stats['normal']
    polluted = stats['polluted']
    defended = stats['defended']
    sent_num = normal + polluted
    packet_size_byte = 1400
    Mbps = 8 * sent_num * packet_size_byte / 1024 / 1024 / (SIMULATION_TIME/TIME_SCALE)

    with open(f"stats.csv", "a") as f:
        f.write(f"{stats['normal']},{stats['polluted']},{stats['defended']},{Mbps}\n")
        
    print(stats)
    print(f"link speed: {Mbps}")
    
if __name__ == "__main__":
    with open("stats.csv", 'w') as f:
        f.write(f"normal, polluted, defended, Mbps\n")
    
    from multiprocessing import Process
    num_samples = 20
    # num_samples = 4
    process_list  = []
    for i in range(num_samples):
        p = Process(target=test)
        p.start()
        process_list.append(p)
    
    for p in process_list:
        p.join()
        